﻿#include "MemoryPool.h"

#include "Render/VulkanContext.h"


CVulkanMemoryChunk::CVulkanMemoryChunk(
    uint32_t InMemoryTypeIndex
    , vk::Device InDevice
    , int64_t InTotalSize
    , int64_t InBlockSize
    ,bool bAutoMap
    )
    :TotalSize(InTotalSize)
    ,BlockSize(InBlockSize)
{
   
    vk::MemoryAllocateInfo MemoryAllocateInfo;
    MemoryAllocateInfo.allocationSize = InTotalSize;
    MemoryAllocateInfo.memoryTypeIndex =InMemoryTypeIndex;
    Memory = CVulkanContext::Get()->GetDevice().allocateMemory(MemoryAllocateInfo);

    if(bAutoMap)
    {
        MappedMemory=(uint8_t*)CVulkanContext::Get()->GetDevice().mapMemory(Memory, 0, TotalSize);
    }

    //填充UnusedBlockOffsets
    for (int64_t i = 0; i < TotalSize / InBlockSize; i++)
    {
        UnusedBlockOffsets.push_back(i * InBlockSize);
    }
}

CVulkanMemoryChunk::~CVulkanMemoryChunk()
{
    CVulkanContext::Get()->GetDevice().freeMemory(Memory);
}


CMemoryPool::CMemoryPool(
    EVulkanMemoryUsage InUsage
    , const SMemoryPoolSizes & InSizes
    ,uint32_t InMemoryTypeIndex)
    : Sizes(InSizes)
    , Usage(InUsage)
    , MemoryTypeIndex(InMemoryTypeIndex)
{
    //初始化Chunks
    int64_t CurrentSize=Sizes.MinBlockSize;
    while(CurrentSize<=Sizes.MaxBlockSize)
    {
        Chunks.Add(TVector<CVulkanMemoryChunk*>());
        CurrentSize*=2;
    }

}

CMemoryPool::~CMemoryPool()
{
    for(auto& Chunks2:Chunks)
    {
        for(auto& Chunk:Chunks2)
        {
            if(Chunk)
            {
                delete Chunk;
            }
        }
    }
    Chunks.Empty();
}

SVulkanMemoryBlock CMemoryPool::AllocateBlock(const vk::MemoryRequirements& Requirement)
{
    assert(Requirement.size<=Sizes.MaxBlockSize&&"需要分配的内存大小超过了最大值，需要调整MaxShaderUniforBufferBlockSize");

    //根据InSize找到合适的chunk组
    //在绑定buffer和memeory时,memory offset必须是内存对齐的整数倍
    //所以块大小不能小于Aligment
    int64_t BlockSize=(uint64_t)Sizes.MinBlockSize;
    int64_t BlockSizeIndex=0;
    {
        while(BlockSize<Requirement.size
            ||BlockSize<Requirement.alignment
        )
        {
            BlockSize<<=1;
            BlockSizeIndex++;
        }
    }

    assert(BlockSize<=Sizes.MaxBlockSize&&"BlockSize超过了最大值，需要调整MaxShaderUniforBufferBlockSize");
    assert(BlockSizeIndex<Chunks.Num()&&"没有找到合适的内存块");

    //查找是否有空闲的大内存块
    int64_t ChunkIndex=-1;
    CVulkanMemoryChunk* Chunk=nullptr;
    if(Chunks[BlockSizeIndex].Num()>0)
    {
        for(int64_t i=0;i<Chunks[BlockSizeIndex].Num();i++)
        {
            if(Chunks[BlockSizeIndex][i] && !Chunks[BlockSizeIndex][i]->IsFull())
            {
                Chunk=Chunks[BlockSizeIndex][i];
                ChunkIndex=i;
                break;
            }
        }
    }

    //没有找到合适的大内存块，创建一个新的
    if(!Chunk)
    {
        //可能存在为空的Chunk
        for(int64_t i=0;i<Chunks[BlockSizeIndex].Num();i++)
        {
            if(!Chunks[BlockSizeIndex][i])
            {
                //这个位置可以用
                ChunkIndex=i;
                break;
            }
        }

        if(ChunkIndex==-1)
        {
            Chunks[BlockSizeIndex].Add(nullptr);
            ChunkIndex=Chunks[BlockSizeIndex].Num()-1;
        }

        Chunk= new CVulkanMemoryChunk(
            MemoryTypeIndex
            ,CVulkanContext::Get()->GetDevice()
            ,Sizes.ChunkSize
            ,BlockSize
            ,Usage==EVulkanMemoryUsage::MappedTexture
            );
        Chunks[BlockSizeIndex][ChunkIndex]=Chunk;
    }
 
    //分配内存
    SVulkanMemoryBlock Block;
    Block.Offset=Chunk->AllocateBlock();
    Block.Memory=Chunk->Memory;
    Block.Size=Requirement.size;
    Block.Capacity=BlockSize;
    Block.ChunkIndex= ChunkIndex;
    Block.MemoryTypeIndex=MemoryTypeIndex;
    Block.Usage=Usage;
    Block.BlockSizeIndex=BlockSizeIndex;
    if(Chunk->HasMappedMemory())
    {
        Block.MappedMemory=Chunk->GetMappedMemory()+Block.Offset;
    }

    assert(Block.Offset<Chunk->GetTotalSize()&&"Offset超过了Chunk的大小");

    return Block;

}

void CMemoryPool::DeallocateBlock(const SVulkanMemoryBlock &Block)
{
    //根据InSize找到合适的chunk组
    int64_t BlockSizeIndex=Block.BlockSizeIndex;
   
    assert(BlockSizeIndex<Chunks.Num()&&"没有找到合适的内存块");

    assert(Block.ChunkIndex<Chunks[BlockSizeIndex].Num()&&"ChunkIndex越界");

    auto Chunk=Chunks[BlockSizeIndex][Block.ChunkIndex];
    assert(Chunk&&"找不到Chunk，可能已经被释放了");
    Chunk->DeallocateBlock(Block);

    //如果Chunk为空，删除Chunk
    if(Chunk->IsEmpty())
    {
        delete Chunk;
        Chunks[BlockSizeIndex][Block.ChunkIndex]=nullptr;
    }


}


void CMemoryPools::Init()
{
    const SMemoryPoolSizes UniformBufferSizes=SMemoryPoolSizes(
        1024*1024*16 //每个大块
        ,16       //每个小块的下限
        ,1024*1024*16  //每个小块的上限
    );

    //顶点缓冲区的内存池尺寸
    //每个顶点要8个float,也就是32字节
    const SMemoryPoolSizes VertexBufferSizes=SMemoryPoolSizes(
        32* (1<< 22)     //每个大块, 支持4M个顶点
        ,32* 16          //每个小块的下限 ，16个顶点
        ,32* (1<< 20)    //每个小块的上限, 1M个顶点
    );

    //索引缓冲区的内存池尺寸
    //每个索引要4个字节(uint32_t)
    const SMemoryPoolSizes IndexBufferSizes=SMemoryPoolSizes(
        4* (1<< 22)     //每个大块, 支持4M个索引
        ,4* 16          //每个小块的下限 ，16个索引
        ,4* (1<< 20)    //每个小块的上限, 1M个索引
    );


    //纹理的内存池尺寸,
    const SMemoryPoolSizes TextureSizes=SMemoryPoolSizes(
        7680* 4320 * 4 *4 //每个大块,4 张16k纹理, 32位RGBA , 一共是512M
        ,4096         //64*64*1*1=4096
        ,7680* 4320 * 4 *4 //每个小块的上限,与大块一样大
    );

    PoolSizes=
    {
        {EVulkanMemoryUsage::UniformBuffer,UniformBufferSizes},
        {EVulkanMemoryUsage::VertexBuffer,VertexBufferSizes},
        {EVulkanMemoryUsage::IndexBuffer,IndexBufferSizes},
        {EVulkanMemoryUsage::Texture,TextureSizes},
        {EVulkanMemoryUsage::MappedTexture,TextureSizes}

    };
}

SVulkanMemoryBlock CMemoryPools::AllocateBlock(
    EVulkanMemoryUsage InUsage
    ,uint32_t InMemoryTypeIndex
    ,const vk::MemoryRequirements& Requirement
    )
{
    SMemoryPoolKey Key;
    Key.MemoryTypeIndex=InMemoryTypeIndex;
    Key.Usage=InUsage;

    auto Found=MemoryPools.Find(Key);
    if(Found)
    {
        return (*Found)->AllocateBlock(Requirement);
    }

    //没有找到对应的内存池，创建一个新的
    MemoryPools.Add(Key,std::make_unique<CMemoryPool>(InUsage,PoolSizes[InUsage],InMemoryTypeIndex));
    return MemoryPools[Key]->AllocateBlock(Requirement);
}

void CMemoryPools::DeallocateBlock(const SVulkanMemoryBlock &Block)
{
    SMemoryPoolKey Key;
    Key.MemoryTypeIndex=Block.MemoryTypeIndex;
    Key.Usage=Block.Usage;

    assert(MemoryPools.Contains(Key)&&"找不到对应的内存池");
    MemoryPools[Key]->DeallocateBlock(Block);
}
